package main

import (
	"io"
	"log"
	"net"
	"net/http"
	"os"
	"time"

	"github.com/gorilla/websocket"
)

// target => client 转发任务
func target_to_client(client *websocket.Conn, target net.Conn, wsType int) {
	// 退出时关闭 websocket 连接
	defer client.Close()
	// 退出时关闭到目标主机的连接
	defer target.Close()
	// target 读超时时间
	readLazyTime := 100 * time.Millisecond

	// 分配 client 下一帧的写缓冲区
	var buf [1024]byte

	for {
		// 设置 target 读超时时间
		target.SetReadDeadline(time.Now().Add(readLazyTime))
		// 从 target 中读取数据，直到超时
		n, err := target.Read(buf[:])
		// os.ErrDeadlineExceeded 在 go1.14 linux 版本中不存在，需要改为如下方式判断超时
		//if err != nil && !errors.Is(err, os.ErrDeadlineExceeded) {
		if err != nil && !os.IsTimeout(err) {
			log.Println(err)
			return
		}

		// 把读到的数据做为一帧发到 client
		if n > 0 {
			// 取得 client 下一帧的 writer
			w, err := client.NextWriter(wsType)
			if err != nil {
				log.Println(err)
				return
			}
			// 写到 client 下一帧的 writer 缓冲区
			_, err = w.Write(buf[:n])
			if err != nil {
				log.Println(err)
				return
			}
			// 关闭下一帧的 writer，真正地发送数据帧
			err = w.Close()
			if err != nil {
				log.Println(err)
				return
			}
		}
	}
}

// client => target 转发任务
func client_to_target(client *websocket.Conn, target net.Conn, wsType int) {
	// 退出时关闭 websocket 连接
	defer client.Close()
	// 退出时关闭到目标主机的连接
	defer target.Close()

	for {
		messageType, r, err := client.NextReader()
		if err != nil {
			log.Println(err)
			return
		}
		if messageType == websocket.CloseMessage {
			return
		} else if messageType == websocket.TextMessage || messageType == websocket.BinaryMessage {
			w, err := client.NextWriter(messageType)
			if err != nil {
				return
			}
			// 转发到 target
			if _, err := io.Copy(target, r); err != nil {
				return
			}
			// 结束本帧数据转发
			if err := w.Close(); err != nil {
				return
			}
		}
	}
}

// 本示例提供 WebSocket <=> TCP Proxy 功能
func wsktTcpHandler(w http.ResponseWriter, r *http.Request) {
	// 从 URL 参数中提取目标主机地址和端口信息
	log.Printf("Accept client %s => %s", r.RemoteAddr, r.URL)
	query := r.URL.Query()
	host := query.Get("host")
	port := query.Get("port")
	targetAddr := net.JoinHostPort(host, port)
	wsType := websocket.TextMessage
	if query.Get("type") == "binary" {
		wsType = websocket.BinaryMessage
	}

	// 建立到目标主机的连接
	log.Printf("Connect to target %s ...\r\n", targetAddr)
	target, err := net.Dial("tcp", targetAddr)
	if err != nil {
		log.Println(err)
		return
	}

	// 升级当前 HTTP 连接为 WebSocket 连接
	log.Println("Upgrade HTTP => WebSocket ...")
	client, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Println("Upgrade failed:", err)
		target.Close()
		return
	}

	// 启动 client 连接服务任务
	log.Printf("Proxy: %s <=> (%s = %s) <=> %s\r\n", client.RemoteAddr(), client.LocalAddr(), target.LocalAddr(), target.RemoteAddr())
	go client_to_target(client, target, wsType)
	go target_to_client(client, target, wsType)
}
